package cas.ihep.spark.rdd;

import cas.ihep.spark.context.JSparkIterationListener;
import cas.ihep.spark.loop.ClientIterativeRunnable;
import cas.ihep.spark.loop.IterationClient;
import cas.ihep.spark.loop.IterationData;
import cas.ihep.spark.loop.IterativeRunnable;
import cas.ihep.spark.util.GlobalConfiguration;
import cas.ihep.spark.util.IterationClientListener;
import cas.ihep.spark.util.ManagedResource;
import cas.ihep.spark.util.NamedResourceManager;
import cas.ihep.spark.util.function.Function2;
import cas.ihep.spark.util.function.Function3;
import org.apache.spark.Partition;
import org.apache.spark.SparkContext;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.broadcast.Broadcast;
import org.apache.spark.rdd.RDD;
import scala.collection.JavaConversions;
import scala.reflect.ClassManifestFactory;
import scala.reflect.ClassTag;

import java.io.Serializable;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.util.List;

public class JIterativeRDD<T> extends JavaRDD<T> {

    private JIterativeRDD(RDD<T> rdd, ClassTag<T> classTag) {
        super(rdd, classTag);
    }

    public static <T> JIterativeRDD<T> from(JavaRDD<T> rdd){
        if(rdd instanceof JIterativeRDD){
            return (JIterativeRDD<T>)rdd;
        }
        return new JIterativeRDD<>(rdd.rdd(),rdd.classTag());
    }

    public List<String> getPreferredLocations(Partition p){
        return JavaConversions.seqAsJavaList(rdd().getPreferredLocations(p));
    }

    public static <T> JIterativeRDD<T> from(RDD<T> rdd){
        return new JIterativeRDD<>(rdd,rdd.elementClassTag());
    }

    public <Send extends Serializable, Receive extends Serializable> void iterate(IterativeRunnable<Send, Receive> iRun,
                                                             IterationClientListener<T, Send, Receive> cRun, Class<Send> cls, Receive zero, Function2<Receive, Receive, Receive> agg){
        IterativeRDDUtil.iterate(rdd(),context(),getNumPartitions(),iRun,cRun,cls,zero,agg);
    }

    @Deprecated
    public <Send extends Serializable, Receive extends Serializable> void iterate(IterativeRunnable<Send, Receive> iRun,
                                                             Function3<T, Send, NamedResourceManager, Receive> cRun,
                                                             Class<Send> cls, Receive zero, Function2<Receive, Receive, Receive> agg){
        final SparkContext sc=context();
        IterativeBroadcastData<T, Send, Receive> param;
        try {
            param=new IterativeBroadcastData<>(new InetSocketAddress(InetAddress.getLocalHost(), 10000),cls,
                    zero,agg,cRun);
        }catch (UnknownHostException error){
            throw new RuntimeException("Should never happens",error);
        }
        param.properties=GlobalConfiguration.properties();
        sc.addSparkListener(new JSparkIterationListener<>(sc, iRun, getNumPartitions(), param.serverAddr));
        final Broadcast<IterativeBroadcastData<T, Send, Receive>> bcast=sc.broadcast(param,ClassManifestFactory.classType(param.getClass()));
        rdd().foreach(new IterativeRDDUtil<T>() {
            @Override
            public void call(T t) throws Exception {
                IterativeBroadcastData<T, Send, Receive> param1=bcast.getValue();
                initExecutorProperties(param1.properties);
                try(NamedResourceManagerImpl closer=new NamedResourceManagerImpl()){
                    InetSocketAddress localAddr=new InetSocketAddress(InetAddress.getLocalHost(),8999);
                    @SuppressWarnings("unchecked")
                    IterationClient<Send, Receive> client1=IterationClient.getOrCreateSingleInstance(BaseBroadcastData.CLIENT_KEY,localAddr,
                            param1.serverAddr,(Class<Send>)param1.sendClassTag.runtimeClass());
                    //noinspection Convert2Lambda
                    client1.waitingForData(new ClientIterativeRunnable<Send, Receive>() {

                        public void run(IterationData<Send, Receive> iData) {
                            iData.output(param1.clientFunc.call(t,iData.input(),closer));
                        }
                    },param1.zero,param1.aggregateFunc);
                }
            }
        });
    }

    @Deprecated
    public <Send extends Serializable, Receive extends Serializable> void iterate(IterativeRunnable<Send, Receive> iterativeRunnable,
                                                             final Function2<ManagedResource<T>, Send, Receive> client, Class<Send> pClass, Receive zero,
                                                             Function2<Receive, Receive, Receive> aggregateFunc){
        final SparkContext sc=context();
        LoopBroadcastParam<T, Send, Receive> param;
        try {
            param=new LoopBroadcastParam<>(new InetSocketAddress(InetAddress.getLocalHost(), 10000),
                    zero,aggregateFunc,client,ClassManifestFactory.classType(pClass));
        }catch (UnknownHostException error){
            throw new RuntimeException("Should never happens",error);
        }
        param.props= GlobalConfiguration.properties();
        sc.addSparkListener(new JSparkIterationListener<>(sc, iterativeRunnable, getNumPartitions(), param.serverAddr));
        final Broadcast<LoopBroadcastParam<T, Send, Receive>> globalConfigBcast=sc.broadcast(param, ClassManifestFactory.classType(param.getClass()));
        rdd().foreach(new IterativeRDDUtil<T>() {
            @Override
            public void call(T t) throws Exception {
                LoopBroadcastParam<T, Send, Receive> param1 = globalConfigBcast.getValue();
                initExecutorProperties(param1.props);
                try (DefaultManagedResource<T> resource = new DefaultManagedResource<>(t)){
                    InetSocketAddress localAddr = new InetSocketAddress(InetAddress.getLocalHost(), 8999);
                    @SuppressWarnings("unchecked")
                    IterationClient<Send, Receive> client1 = IterationClient.getOrCreateSingleInstance(BaseBroadcastData.CLIENT_KEY,
                            localAddr, param1.serverAddr, (Class<Send>) param1.classTag.runtimeClass());
                    client1.waitingForData(client11 -> client11.output(param1.client.call(resource, client11.input())), param1.zero, param1.aggregate);
                }
            }
        });
    }

}
